from cio.jsonl import JsonLReader, JsonLWriter
from config import Config
from tqdm import tqdm
from utils.splitter import code_split, nl_split
import pickle
import string

file = Config.BASE_DIR / 'resource' / 'data' / 'consum' / 'context.jsonl'
drop_list = set()
reader = JsonLReader(file)


def drop_item(item: dict) -> bool:
    # 大于500个token
    try:
        if len(code_split(item['method'])) > 500:
            return True
        # 摘要3~30个单词
        if not 3 <= len(nl_split(item['summary'])) <= 28:  # 要加一个eos,bos,长度是30
            # 摘要里面包含非ascii的
            return True
        for char in item['summary']:
            # 不允许出现非ascii
            if char not in string.printable:
                return True
        if len(code_split(item['package_name'])) + \
                len(code_split(item['class_signature'])) + \
                sum([len(code_split(each)) for each in item['called_statements']]) + \
                sum([len(code_split(each)) for each in item['callee_signatures']]) > 500:
            return True
        return False
    except:
        return True


for index, data in enumerate(tqdm(list(reader.yield_read()))):
    if drop_item(data):
        drop_list.add(index)

pickle.dump(drop_list, open(Config.BASE_DIR / 'resource' / 'data' / 'consum' / 'drop_list.pkl', 'wb'))
